
#ifndef DIRECT_MATRIX_HPP_INCLUDED
#define DIRECT_MATRIX_HPP_INCLUDED

#include <aftermath/numerics/matrix.hpp>

#include <algorithm>
#include <cstring>

namespace aftermath
{
    namespace numerics
    {
        template <int M, int N, typename T>
        struct matrix<M, N, T, void>
        {
            template <int K, int L, typename U, typename S>
            friend struct matrix;

            typedef T data_type;

            static const int height = M;
            static const int width = N;
            
            static const bool is_square = (M == N);
            static const int memory_wasted = 0;
        
            static const int required_align = 1;
            static const int stride = N;
            static const int required_memory = M * N * sizeof(data_type);

            matrix()
                : m_data(nullptr)
            {
                this->m_data = reinterpret_cast<data_type*>(malloc(required_memory));
            }

            matrix(const matrix& other) 
                : m_data(nullptr)
            {
                this->m_data = reinterpret_cast<data_type*>(malloc(required_memory));
                std::memcpy(this->m_data, other.m_data, M * N * sizeof(data_type)); 
            }

            matrix(matrix&& other)
                : m_data(nullptr)
            {
                *this = std::move(other);
            }

            matrix& operator =(const matrix& other) 
            {
                if (this != &other) std::memcpy(this->m_data, other.m_data, required_memory);
                return *this;
            }

            matrix& operator =(matrix&& other)
            {
                if (this != &other)
                {
                    free(this->m_data);
                    this->m_data = other.m_data;
                    other.m_data = nullptr;
                }
                return *this;
            }

            ~matrix()
            {
                free(this->m_data);
                this->m_data = nullptr;
            }
            
            void fill(data_type value) { for (int i = 0; i < M; i++) for (int j = 0; j < N; j++) this->m_data[i * stride + j] = value; }
            
            void fill_row(int row, data_type value) { for (int j = 0; j < N; j++) this->m_data[row * stride + j] = value; }

            void clear() { for (int i = 0; i < M; i++) for (int j = 0; j < N; j++) this->m_data[i * stride + j] = 0; }

            void clear_row(int row) { for (int j = 0; j < N; j++) this->m_data[row * stride + j] = 0; }

            data_type get_element(int row, int column) const { return this->m_data[row * stride + column]; }

            template <int row, int column>
            void set_element(data_type value) { this->m_data[row * stride + column] = value; }

            void set_element(int row, int column, data_type value) { this->m_data[row * stride + column] = value; }
            
            template <int row, int column>
            void clear_element() { this->m_data[row * stride + column] = 0; }
            
            void clear_element(int row, int column) { this->m_data[row * stride + column] = 0; }

            void transpose_self();
            matrix<N, M, T, void> transpose() const;

            template <int K>
            matrix<M, K, T, void> multiply(const matrix<N, K, T, void>& other) const;

            template <int K>
            matrix<M, K, T, void> fast_multiply(const matrix<N, K, T, void>& other) const;
            
            matrix<M, N, T, void> operator +(const matrix<M, N, T, void>& other)
            {
                matrix<M, N, T, void> result;
                for (int i = 0; i < M; i++) for (int j = 0; j < N; j++) result.m_data[i * stride + j] = this->m_data[i * stride + j] + other.m_data[i * stride + j];
                return result;
            }
            
            matrix<M, N, T, void> operator -(const matrix<M, N, T, void>& other)
            {
                matrix<M, N, T, void> result;
                for (int i = 0; i < M; i++) for (int j = 0; j < N; j++) result.m_data[i * stride + j] = this->m_data[i * stride + j] - other.m_data[i * stride + j];
                return result;
            }

            matrix<M, N, T, void> operator *(const matrix<M, N, T, void>& other)
            {
                matrix<M, N, T, void> result;
                for (int i = 0; i < M; i++) for (int j = 0; j < N; j++) result.m_data[i * stride + j] = this->m_data[i * stride + j] * other.m_data[i * stride + j];
                return result;
            }

            matrix<M, N, T, void> operator /(const matrix<M, N, T, void>& other)
            {
                matrix<M, N, T, void> result;
                for (int i = 0; i < M; i++) for (int j = 0; j < N; j++) result.m_data[i * stride + j] = this->m_data[i * stride + j] / other.m_data[i * stride + j];
                return result;
            }

            matrix<M, N, T, void>& operator +=(const matrix<M, N, T, void>& other)
            {
                for (int i = 0; i < M; i++) for (int j = 0; j < N; j++) this->m_data[i * stride + j] += other.m_data[i * stride + j];
                return *this;
            }

            matrix<M, N, T, void>& operator -=(const matrix<M, N, T, void>& other)
            {
                for (int i = 0; i < M; i++) for (int j = 0; j < N; j++) this->m_data[i * stride + j] -= other.m_data[i * stride + j];
                return *this;
            }

            matrix<M, N, T, void>& operator *=(const matrix<M, N, T, void>& other)
            {
                for (int i = 0; i < M; i++) for (int j = 0; j < N; j++) this->m_data[i * stride + j] *= other.m_data[i * stride + j];
                return *this;
            }

            matrix<M, N, T, void>& operator /=(const matrix<M, N, T, void>& other)
            {
                for (int i = 0; i < M; i++) for (int j = 0; j < N; j++) this->m_data[i * stride + j] /= other.m_data[i * stride + j];
                return *this;
            }

        private:
            data_type* m_data;
        };
        
        template <int M, int N, typename T>
        void matrix<M, N, T, void>::transpose_self()
        {
            static_assert(M == N, "transpose_self only works on square matrices.");
            for (int i = 0; i < M; i++) 
            {
                for (int j = i + 1; j < N; j++)
                {
                    T temp = this->m_data[j * stride + i];
                    this->m_data[j * stride + i] = this->m_data[i * stride + j];
                    this->m_data[i * stride + j] = temp;
                };
            }
        }
    
        template <int M, int N, typename T>
        matrix<N, M, T, void> matrix<M, N, T, void>::transpose() const
        {
            typedef matrix<M, N, T, void> this_type;
            typedef matrix<N, M, T, void> result_type;

            matrix<N, M, T, void> result;
            for (int i = 0; i < M; i++) for (int j = 0; j < N; j++) result.m_data[j * result_type::stride + i] = this->m_data[i * this_type::stride + j];
            return result;
        }
    
        template <int M, int N, typename T>
        template <int K>
        matrix<M, K, T, void> matrix<M, N, T, void>::multiply(const matrix<N, K, T, void>& other) const
        {
            typedef matrix<M, N, T, void> left_type;
            typedef matrix<N, K, T, void> right_type;
            typedef matrix<M, K, T, void> result_type;

            matrix<M, K, T, void> result;
            for (int i = 0; i < M; i++)
            {
                for (int j = 0; j < K; j++)
                {
                    T sum = 0;
                    for (int l = 0; l < N; l++) sum += this->m_data[i * left_type::stride + l] * other.m_data[l * right_type::stride + j];
                    result.m_data[i * result_type::stride + j] = sum;
                }
            }
            return result;
        }
    }
}


#endif // DIRECT_MATRIX_HPP_INCLUDED
